import os, sys
from tqdm import trange
import numpy as np
from astropy.io import fits
from astropy.table import Table
from scipy.optimize import curve_fit
from scipy.stats import sigmaclip
from scipy.ndimage import binary_dilation
from matplotlib import pyplot as plt


def flag_bias(image):
    index = np.logical_or(np.abs(image) > 50, np.isnan(image))
    index = binary_dilation(index, iterations=3)
    clip = sigmaclip(image[~index], 5, 5)[0]
    mnn, sigma = clip.mean(), clip.std()
    flagged = np.abs(image - mnn) > sigma * 3

    index2 = np.logical_or(binary_dilation(index) * flagged, index)
    while np.equal(index, index2) is False:
        index = index2
        index2 = np.logical_or(binary_dilation(index) * flagged, index)
    return ~index


def bias_stats(flist, fbias, fnoise, fgain, fout, foutlier, imgconfig, fmask=None, ffig=None, check=False):
    # load reference images
    bias = fits.getdata(fbias)
    gainmap = fits.getdata(fgain)
    mask = np.isfinite(gainmap)
    if fmask is not None:
        mask = mask * (fits.getdata(fmask) < 32)

    # calculate noise from noise map
    with fits.open(fnoise) as img:
        data_stack = img[0].data
        noise_stack = sigmaclip(data_stack[mask], 5, 5)[0].mean()
    # calculate fraction and plot histogram
    n = len(flist)
    name = []
    noise = np.zeros(n,dtype=np.float32)
    noise_channels = np.zeros(imgconfig.nchan)
    npix15 = np.zeros(n,dtype=np.int32)
    npix25 = np.zeros(n, dtype=np.int32)
    frac15 = np.zeros(n, dtype=np.float32)
    frac25 = np.zeros(n, dtype=np.float32)
    ########for CMOS channel stat#########
    npix15_channels = np.zeros((imgconfig.nchan,n),dtype=np.int32)
    frac15_channels = np.zeros((imgconfig.nchan,n),dtype=np.float32)
    npix25_channels = np.zeros((imgconfig.nchan,n),dtype=np.int32)
    frac25_channels = np.zeros((imgconfig.nchan,n),dtype=np.float32)
    #####################################
    ncol = 4
    nrow = n // ncol
    if n > nrow * ncol:
        nrow += 1
    plt.figure(figsize=(ncol * 4, nrow * 4))
    for j in range(imgconfig.nchan):
        data_stack_channel = data_stack[imgconfig.mos_y0[j]:imgconfig.mos_y0[j]+imgconfig.ny[j],imgconfig.mos_x0[j]:imgconfig.mos_x0[j]+imgconfig.nx[j]]
        mask_channel = mask[imgconfig.mos_y0[j]:imgconfig.mos_y0[j]+imgconfig.ny[j],imgconfig.mos_x0[j]:imgconfig.mos_x0[j]+imgconfig.nx[j]]
        noise_channels[j] = sigmaclip(data_stack_channel[mask_channel], 4, 4)[0].mean()
    
    for i in trange(n, desc='processing bias'):
        name.append(os.path.basename(flist[i]))
        with fits.open(flist[i]) as img:
            data = (img[0].data - bias) * gainmap
            index = flag_bias(data) * mask
            fits.writeto(foutlier+"bias_stat_flag_"+str(i+1)+".fits",data=index.astype('uint8'),overwrite=True)
            noise[i] = sigmaclip(data[index], 4, 4)[0].std()
            
            #print(noise_channels[i,:])
#"""
            index1 = (np.abs(data) > 15) * index
            fits.writeto(foutlier+"bias_outlier_15_"+str(i+1)+".fits",data = index1.astype('uint8'),overwrite=True)
            npix15[i] = np.sum(index1)
            frac15[i] = npix15[i] / np.sum(index)
            index2 = (np.abs(data) > 25) * index
            fits.writeto(foutlier+"bias_outlier_25_"+str(i+1)+".fits",data = index2.astype('uint8'),overwrite=True)
            npix25[i] = np.sum(index2)
            frac25[i] = npix25[i] / np.sum(index)
#"""         
            ############for CMOS channel stats##########
            for j in range(imgconfig.nchan):
                npix15_channels[j,i] = np.sum(index1[imgconfig.mos_y0[j]:imgconfig.mos_y0[j]+imgconfig.ny[j],imgconfig.mos_x0[j]:imgconfig.mos_x0[j]+imgconfig.nx[j]])
                frac15_channels[j,i] = npix15_channels[j,i]/np.sum(index[imgconfig.mos_y0[j]:imgconfig.mos_y0[j]+imgconfig.ny[j],imgconfig.mos_x0[j]:imgconfig.mos_x0[j]+imgconfig.nx[j]])
                npix25_channels[j,i] = np.sum(index2[imgconfig.mos_y0[j]:imgconfig.mos_y0[j]+imgconfig.ny[j],imgconfig.mos_x0[j]:imgconfig.mos_x0[j]+imgconfig.nx[j]])
                frac25_channels[j,i] = npix25_channels[j,i]/np.sum(index[imgconfig.mos_y0[j]:imgconfig.mos_y0[j]+imgconfig.ny[j],imgconfig.mos_x0[j]:imgconfig.mos_x0[j]+imgconfig.nx[j]])
            ###########################################
            if check:
                hdu = fits.PrimaryHDU(index1.astype(np.int16) + index2.astype(np.int16) * 2)
                hdu2 = fits.ImageHDU(data)
                hdulist = fits.HDUList([hdu, hdu2])
                dname = os.path.dirname(flist[i])
                fname = os.path.basename(flist[i]).replace('bias-', 'bmask-')
                hdulist.writeto('{}/{}'.format(dname, fname), overwrite=True)

            resolution = noise[i] / 5
            bins = (np.arange(161) - 80) * resolution
            plt.subplot(nrow, ncol, i+1)
            plt.hist(data[mask], bins=bins, color='lightgreen', log=True)
            plt.xlim([bins[0], bins[-1]])
            plt.xlabel('Offset (e-)')
            plt.ylabel('Number of pixels')
            plt.title(os.path.basename(flist[i]))

    plt.tight_layout()
    if ffig is None:
        ffig = fout.replace('.tab', '.png')
    plt.savefig(ffig)
    plt.close()
    
    # stats for all
    name = np.append(name, ['median', 'max'])
    noise = np.append(noise, [np.median(noise), np.max(noise)])
    npix15 = np.append(npix15, [np.median(npix15), np.max(npix15)])
    frac15 = np.append(frac15, [np.median(frac15), np.max(frac15)])
    npix25 = np.append(npix25, [np.median(npix25), np.max(npix25)])
    frac25 = np.append(frac25, [np.median(frac25), np.max(frac25)])
    
    # output
    tab = Table()
    tab['name'] = name
    tab['noise_indiv'] = noise
    tab['noise_indiv'].unit = 'ph'
    tab['noise'] = np.float32(noise_stack)
    tab['noise'].unit = 'ph'
    tab['npix15'] = np.array(npix15).astype(int)
    tab['frac15'] = frac15
    tab['npix25'] = np.array(npix25).astype(int)
    tab['frac25'] = frac25
    tab.sort('name')
    tab.write(fout, format='ipac', overwrite=True)

    tab_channel = Table()
    tab_channel['channel'] = np.append(np.arange(imgconfig.nchan).astype(str),np.array(['median','max']))
    tab_channel['noise'] = np.append(noise_channels,np.array([np.median(noise_channels),np.max(noise_channels)]))
    tab_channel['noise'].unit = 'ph'
    ##############for CMOS channel stat############
    npix15_chan = np.median(npix15_channels,axis=1)
    frac15_chan = np.median(frac15_channels,axis=1)
    npix25_chan = np.median(npix25_channels,axis=1)
    frac25_chan = np.median(frac25_channels,axis=1)
    tab_channel['npix15'] = np.append(npix15_chan,np.array([np.median(npix15_chan),np.max(npix15_chan)]))
    tab_channel['frac15'] = np.append(frac15_chan,np.array([np.median(frac15_chan),np.max(frac15_chan)]))
    tab_channel['npix25'] = np.append(npix25_chan,np.array([np.median(npix25_chan),np.max(npix25_chan)]))
    tab_channel['frac_25'] = np.append(frac25_chan,np.array([np.median(frac25_chan),np.max(frac25_chan)]))
    ###############################################
    tab_channel.write(fout[:-4]+"_chan.tab", format='ipac', overwrite=True)

def dark_stats(flist, fout, t0, fmask=None, fgain=None, foutlier=None):
    # load reference images
    print("dark stat fmask:",fmask)
    if fgain is None:
        gainmap = 1
        mask = np.isfinite(fits.getdata(flist[0]))
    else:
        gainmap = fits.getdata(fgain)
        mask = np.isfinite(gainmap)
    if fmask is not None:
        mask = mask * (fits.getdata(fmask) < 32)

    # calculate fraction
    n = len(flist)
    name = []
    tsys = np.zeros(n, dtype=np.float32)
    dc = np.zeros(n, dtype=np.float32)
    noise = np.zeros(n, dtype=np.float32)
    
    npix08 = np.zeros(n, dtype=np.int32)
    frac08 = np.zeros(n, dtype=np.float32)
    npix16 = np.zeros(n, dtype=np.int32)
    frac16 = np.zeros(n, dtype=np.float32)

    for i in trange(n, desc='processing dark'):
        name.append(os.path.basename(flist[i]))
        with fits.open(flist[i]) as img:
            data = img[0].data * gainmap
            tsys[i] = img[0].header['SYSTEMP']
            #print("mask_sum",np.sum(mask))
            index = (data > 0.08) * mask
            npix08[i] = np.sum(index)
            frac08[i] = npix08[i] / np.sum(mask)
            fits.writeto(foutlier+flist[i].split('/')[-1][-9:-5]+"_08.fits",data=index.astype('uint8'),overwrite=True)
            
            index = (data > 0.16) * mask
            npix16[i] = np.sum(index)
            frac16[i] = npix16[i] / np.sum(mask)
            fits.writeto(foutlier+flist[i].split('/')[-1][-9:-5]+"_16.fits",data=index.astype('uint8'),overwrite=True)
            clip = sigmaclip(data[mask], 5, 5)[0]
            dc[i] = clip.mean()
            noise[i] = clip.std()

    # stats for all
    index = (tsys >= t0 - 10 - 1e-5) * (tsys <= t0 + 10 + 1e-5)
    index2 = (tsys >= t0 - 10 - 1e-5) * (tsys <= t0 + 1e-5)
    name = np.append(name, ['median', 'max'])
    tsys = np.append(tsys, [np.nan, np.nan])
    
    npix08 = np.append(npix08, [np.median(npix08[index]), np.max(npix08[index])])
    frac08 = np.append(frac08, [np.median(frac08[index]), np.max(frac08[index])])
    npix16 = np.append(npix16, [np.median(npix16[index]), np.max(npix16[index])])
    frac16 = np.append(frac16, [np.median(frac16[index]), np.max(frac16[index])])

    dc = np.append(dc, [np.median(dc[index2]), np.max(dc[index2])])
    noise = np.append(noise, [np.median(noise[index2]), np.max(noise[index2])])
    
    # output
    tab = Table()
    tab['name'] = name
    tab['tsys'] = tsys
    tab['tsys'].unit = 'K'
    tab['dark'] = dc.astype(np.float32)
    tab['dark'].unit = 'ph/s'
    tab['dark_std'] = noise.astype(np.float32)
    tab['dark_std'].unit = 'ph/s'
    tab['npix08'] = np.array(npix08).astype(int)
    tab['frac08'] = frac08
    tab['npix16'] = np.array(npix16).astype(int)
    tab['frac16'] = frac16
    tab.sort(['tsys', 'name'])
    tab.write(fout, format='ipac', overwrite=True)

def _expfunc(x, x0, a, b, c):
    return np.exp((x - x0) / a) * b + c


def plot_dark_stats(ftab, ffig, t0):
    tab = Table.read(ftab, format='ipac')
    tab = tab[np.isfinite(tab['tsys'])]
    index = (tab['tsys'] >= t0 - 10 - 1e-5) * (tab['tsys'] <= t0 + 10 + 1e-5)
    tab = tab[index]
    n = len(tab)
    plt.figure(figsize=(11, 5))

    # plot dark current
    plt.subplot(121)
    plt.xlabel('Temperature (K)')
    plt.ylabel('Dark Current (e-/s)')
    plt.scatter(tab['tsys'], tab['dark'], color='g')
    ax = plt.gca()
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    x = np.arange(101) / 100 * (xlim[1] - xlim[0]) + xlim[0]
    # noinspection PyBroadException
    try:
        if n > 4:
            popt, pcov = curve_fit(_expfunc, tab['tsys'], tab['dark'],
                                   p0=[100, 4, 1, np.min(tab['tsys'])])
            l1, = plt.plot(x, _expfunc(x, *popt), color='r')
            plt.legend([l1, ], ['exp((t - {:.1f}) / {:.2f}) * {:.1e} + {:.4f}'.format(*popt), ])
    except Exception:
        print('WARNING: unsuccessful fitting')
    plt.plot([t0, t0], ylim, color='b', linestyle='--', zorder=-1)
    plt.xlim(*xlim)
    plt.ylim(*ylim)

    # plot >0.08 fraction
    plt.subplot(122)
    plt.xlabel('Temperature (K)')
    plt.ylabel('Fraction (> 0.08 e-/s)')
    plt.bar(tab['tsys'], tab['frac'], color='g')

    plt.tight_layout()
    plt.savefig(ffig)
    plt.close()


def get_prnu(image, varmap):
    var_tot = sigmaclip(image, 5, 5)[0].var(ddof=1)  # don't use the clipped data
    var_pix = np.median(varmap)
    print("var_tot",var_tot,"var_pix",var_pix,"image_mean",image.mean())
    return np.sqrt(var_tot - var_pix) / image.mean()


def flat_stats(flist, fout, fmask=None, fgain=None, band=None, modlist=None):
    if modlist is not None:
        if len(flist) != len(modlist):
            raise Exception('model list should have the same length as the file list')
    # load reference images
    print("fgain",fgain)
    if fgain is None:
        gainmap = 1
        print("flist[0]",flist[0])
        mask = np.isfinite(fits.getdata(flist[0]))
    else:
        gainmap = fits.getdata(fgain)
        print("gainmap_shape",gainmap.shape)
        mask = np.isfinite(gainmap)
    if fmask is not None:
        mask = mask * (fits.getdata(fmask,dtype='>i2') < 32)

    # calculate fraction
    n = len(flist)
    name = []
    wave = np.zeros(n, dtype=np.float32)
    flux = np.zeros(n, dtype=np.float32)
    prnu = np.zeros(n, dtype=np.float32)
    prnu_corr = np.full(n, np.nan, dtype=np.float32)
    gains = []
    
    #region = [int(4085-(6024//2)),int(4085+(6024//2)),int(4716-8011//2),int(4716+8011//2)] #add
    
    for i in trange(n, desc='processing flat'):
        name.append(os.path.basename(flist[i]))
        print("flist[i]",flist[i])
        with fits.open(flist[i]) as img:
            wave[i] = img[0].header['WAVELEN']
            if 'GAIN' in img[0].header:
                gains.append(img[0].header['GAIN'])

#            data = img[0].data * gainmap
#            varmap = (img[1].data * gainmap) ** 2
#            flux[i] = np.median(data[mask])
#            prnu[i] = get_prnu(data[mask], varmap[mask]) * 100
            
            #mask = mask[region[2]:region[3],region[0]:region[1]] #add
            
            data = img[0].data * gainmap
            #data = data[region[2]:region[3],region[0]:region[1]] #add
            varmap = (img[1].data * gainmap) ** 2
            #varmap = varmap[region[2]:region[3],region[0]:region[1]] #add
            flux[i] = np.median(data[mask])
            prnu[i] = get_prnu(data[mask], varmap[mask]) * 100 

            if modlist is not None:
                mod = fits.getdata(modlist[i])
                mod = mod / np.mean(mod)
                
                #mod = mod[region[2]:region[3],region[0]:region[1]]
                
                data_corr = data / mod
                varmap_corr = varmap / (mod ** 2)
                prnu_corr[i] = get_prnu(data_corr[mask], varmap_corr[mask]) * 100
                fits.writeto(fout.replace(".tab","_var.fits"),data=varmap_corr,overwrite=True)
                fits.writeto(fout.replace(".tab","_corr.fits"),data=data_corr,overwrite=True)
    # output
    tab = Table()
    tab['name'] = name
    tab['wave'] = wave
    tab['wave'].unit = 'nm'
    tab['prnu'] = prnu
    tab['prnu'].unit = '%'
    tab['prnu_corr'] = prnu_corr
    tab['prnu_corr'].unit = '%'
    tab['flux'] = flux
    if fgain is None:
        header = fits.getheader(flist[0])
        if 'BUNIT' in header:
            tab['flux'].unit = header['BUNIT']
        else:
            tab['flux'].unit = 'adu'
    else:
        tab['flux'].unit = 'ph'
    if len(gains) == len(tab):
        tab['gain'] = np.array(gains).astype('float32')
        tab['gain'].unit = 'ph/adu'
    tab.sort(['wave', 'name'])
    if band is not None:
        tab.meta['keywords'] = {'band': {'value': band}}
    tab.write(fout, format='ipac', overwrite=True)
